"""
Tensorflow 2.x
"""
import tensorflow as tf


def check_modle_in_tensorboard(model, logdir, x):

    @tf.function
    def traceme(x):
        return model(x)

    writer = tf.summary.create_file_writer(logdir)
    tf.summary.trace_on(graph=True, profiler=False)
    # Forward pass
    traceme(x)
    with writer.as_default():
        tf.summary.trace_export(name="model_trace", step=0)
    print(f'Modle has already been written into log dir ' + logdir)
